#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project : aimix
@File    : extract_audio.py
@IDE     : PyCharm
@Author  : admin
@Date    : 2025/4/3 13:30
"""
import os
import tempfile
import whisper
from datetime import datetime
from moviepy.editor import VideoFileClip
from utils.tools import format_timestamp
from pydub import AudioSegment
from utils.tools import parallel_run

from log import log

def extract_audio_from_video(input_video, output_folder, suffix='.mp3'):
    """
    从视频中提取音频文件
    """
    video = VideoFileClip(input_video)
    audio = video.audio
    f = os.path.join(output_folder, datetime.now().strftime("%Y%m%d%H%M"))
    os.makedirs(f, exist_ok=True)
    audio_file = os.path.normpath(tempfile.mktemp(suffix=suffix, dir=f))
    audio.write_audiofile(audio_file, logger=None)
    return audio_file


def batch_extract_audio_from_video(input_files, output_folder, max_workers=5):
    """
    批量从视频中提取音频文件
    """
    result = parallel_run(extract_audio_from_video, [(f, output_folder) for f in input_files], max_workers=max_workers)
    return result


def transcribe_audio_to_srt(audio_path: str, output_srt_path: str, model_size: str = "base"):
    """
    使用 Whisper 模型将音频转为字幕文件（.srt 格式）

    参数：
        audio_path (str): 输入音频文件路径（支持 .mp3/.wav/.m4a 等）
        output_srt_path (str): 输出字幕文件路径
        model_size (str): Whisper 模型尺寸，默认 base（可选 tiny, base, small, medium, large）
    """
    model = whisper.load_model(model_size)
    result = model.transcribe(audio_path, fp16=False)  # CPU 设置 fp16=False
    log.debug("导出 SRT 文件...")
    with open(output_srt_path, "w", encoding="utf-8") as f:
        for i, segment in enumerate(result["segments"], 1):
            start = segment["start"]
            end = segment["end"]
            text = segment["text"].strip()

            f.write(f"{i}\n")
            f.write(f"{format_timestamp(start)} --> {format_timestamp(end)}\n")
            f.write(f"{text}\n\n")

    log.debug(f"字幕已保存到 {output_srt_path}")


def merge_mp3_files(input_files, output_file):
    """
    合并多个MP3文件为一个文件

    :param input_files: 输入的MP3文件列表
    :param output_file: 输出的MP3文件路径
    """
    combined = AudioSegment.silent(duration=0)

    # 遍历输入文件列表
    for file in input_files:
        # 加载MP3文件
        audio = AudioSegment.from_mp3(file)
        # 将当前音频段拼接到总音频段
        combined += audio

    # 导出合并后的音频
    combined.export(output_file, format="mp3")
    log.debug(f"合并后的文件已保存到 {output_file}")
    return output_file


def extract_video_and_audio(video_path, output_path):
    """
    从视频中提取视频和音频，并分别保存为单独的文件
    :param video_path: 输入视频的路径
    :param output_path: 输出的路径
    """
    # 加载视频文件

    filename = os.path.basename(video_path).split(".")[0]
    output_path = os.path.normpath(os.path.join(output_path, filename))
    os.makedirs(output_path, exist_ok=True)

    output_video_path = os.path.normpath(os.path.join(output_path, "video.mp4"))
    output_audio_path = os.path.normpath(os.path.join(output_path, "audio.mp3"))

    clip = VideoFileClip(video_path)
    audio = clip.audio
    if audio:
        audio.write_audiofile(output_audio_path, logger=None)
        log.debug(f"音频已保存到 {output_audio_path}")
    else:
        log.debug("视频中没有音频")

    # 保存视频（无音频）
    clip.write_videofile(output_video_path, codec="libx264", audio=False, threads=4, logger=None)
    log.debug(f"视频已保存到 {output_video_path}")


def batch_extract_video_and_audio(input_files, output_folder, max_workers=5):
    """
    批量从视频中提取视频和音频，并分别保存为单独的文件
    :param input_files: 输入视频的路径列表
    :param output_folder: 输出的路径
    :param max_workers: 并发数
    """
    result = parallel_run(extract_video_and_audio, [(f, output_folder) for f in input_files], max_workers=max_workers)
    return result


